#include "udp.h"
#include "ip.h"
#include "icmp.h"

/**
 * @brief udp处理程序表
 * 
 */
map_t udp_table;

/**
 * @brief udp伪校验和计算
 * 
 * @param buf 要计算的包
 * @param src_ip 源ip地址
 * @param dst_ip 目的ip地址
 * @return uint16_t 伪校验和
 */
static uint16_t udp_checksum(buf_t *buf, uint8_t *src_ip, uint8_t *dst_ip)
{
    uint16_t total_len16 = swap16(buf->len);

    // 添加UDP伪头部 
    buf_add_header(buf, sizeof(udp_peso_hdr_t));
    udp_peso_hdr_t* udppt = (udp_peso_hdr_t*)buf->data;

    // 暂存IP报头
    udp_peso_hdr_t temp = *udppt;

    // 填写伪头部字段
    memcpy(udppt->src_ip, src_ip, NET_IP_LEN);
    memcpy(udppt->dst_ip, dst_ip, NET_IP_LEN);
    udppt->placeholder = 0;
    udppt->protocol = NET_PROTOCOL_UDP;
    udppt->total_len16 = total_len16;

    // 计算校验和 
    size_t padding_len = (buf->len % 2);
    buf_add_padding(buf, padding_len);
    uint16_t res = checksum16((uint16_t*)buf->data, buf->len);
    buf_remove_padding(buf, padding_len);

    // 恢复IP报头 去除伪头部
    *udppt = temp;
    buf_remove_header(buf, sizeof(udp_peso_hdr_t));

    return res;
}

/**
 * @brief 处理一个收到的udp数据包
 * 
 * @param buf 要处理的包
 * @param src_ip 源ip地址
 */
void udp_in(buf_t *buf, uint8_t *src_ip)
{
    // 如果数据包长度小于UDP报文头长度，丢弃之
    if(buf->len < sizeof(udp_hdr_t))            return;

    // 如果包长度小于UDP首部长度字段给出的长度，丢弃之
    udp_hdr_t* udph = (udp_hdr_t*) buf->data;
    if(buf->len < swap16(udph->total_len16))    return;

    // 计算校验和
    uint16_t temp = udph->checksum16;
    udph->checksum16 = 0;
    if(temp != udp_checksum(buf, src_ip, net_if_ip))    return;
    udph->checksum16 = temp;
 
    // 交付给对应的端口号的处理函数处理 
    uint16_t dst_port = swap16(udph->dst_port16);
    udp_handler_t* handler = map_get(&udp_table, &dst_port);

    // 如果没有查到对应的处理函数
    if(handler == NULL){
        // 为报文添加IP头并发送ICMP端口不可达差错报文
        buf_add_header(buf, sizeof(ip_hdr_t));

        // TODO：这里其实要填写IP报文头，但是由于我们根本没有实现ICMP端口不可达差错报文的应答，所以不填也行

        icmp_unreachable(buf, src_ip, ICMP_CODE_PORT_UNREACH);
    }
    // 如果能找到
    else{
        // 去除UDP报文并调用handler处理
        buf_remove_header(buf, sizeof(udp_hdr_t));
        (*handler)(buf->data, buf->len, src_ip, dst_port); 
    }

}

/**
 * @brief 处理一个要发送的数据包
 * 
 * @param buf 要处理的包
 * @param src_port 源端口号
 * @param dst_ip 目的ip地址
 * @param dst_port 目的端口号
 */
void udp_out(buf_t *buf, uint16_t src_port, uint8_t *dst_ip, uint16_t dst_port)
{
    // 为报文添加UDP报头
    buf_add_header(buf, sizeof(udp_hdr_t));

    // 填写UDP报头
    udp_hdr_t* udph = (udp_hdr_t*)buf->data;
    udph->src_port16 = swap16(src_port);
    udph->dst_port16 = swap16(dst_port);
    udph->total_len16 = swap16(buf->len);

    // 最后计算校验和  
    udph->checksum16 = 0;
    udph->checksum16 = udp_checksum(buf, net_if_ip, dst_ip);

    // 将报文交付IP层
    ip_out(buf, dst_ip, NET_PROTOCOL_UDP);
}

/**
 * @brief 初始化udp协议
 * 
 */
void udp_init()
{
    map_init(&udp_table, sizeof(uint16_t), sizeof(udp_handler_t), 0, 0, NULL);
    net_add_protocol(NET_PROTOCOL_UDP, udp_in);
}

/**
 * @brief 打开一个udp端口并注册处理程序
 * 
 * @param port 端口号
 * @param handler 处理程序
 * @return int 成功为0，失败为-1
 */
int udp_open(uint16_t port, udp_handler_t handler)
{
    return map_set(&udp_table, &port, &handler);
}

/**
 * @brief 关闭一个udp端口
 * 
 * @param port 端口号
 */
void udp_close(uint16_t port)
{
    map_delete(&udp_table, &port);
}

/**
 * @brief 发送一个udp包
 * 
 * @param data 要发送的数据
 * @param len 数据长度
 * @param src_port 源端口号
 * @param dst_ip 目的ip地址
 * @param dst_port 目的端口号
 */
void udp_send(uint8_t *data, uint16_t len, uint16_t src_port, uint8_t *dst_ip, uint16_t dst_port)
{
    buf_init(&txbuf, len);
    memcpy(txbuf.data, data, len);
    udp_out(&txbuf, src_port, dst_ip, dst_port);
}